import cv2 as cv
import numpy as np
import sys
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # 空间三维画图

#图像unsharp mask 增强
def unsharp_mask(im, k = 1.0, r = 31):
    '''
    im: input image
    k: degree of enhancing
    r: radius
    '''
    blured = cv.GaussianBlur(im, (r,r), 0)
    enhanced = (k + 1) * im - k * blured
    enhanced[enhanced < 0] = 0
    enhanced[enhanced > 255] = 255
    enhanced = enhanced.astype(np.uint8)
    return enhanced

#彩色图自适应二值化
def color2bin(im, blk_size = 41, C = 15):
    '''
    im: input color image
    blk_size:  block size for adaptiveThreshold
    C:   C paramter for adaptiveThreshold
    '''
    gray = cv.cvtColor(im, cv.COLOR_BGR2GRAY)
    bin = cv.adaptiveThreshold(gray, 255, cv.ADAPTIVE_THRESH_MEAN_C, cv.THRESH_BINARY_INV, blk_size, C)
    return bin

#二值化图像膨胀操作
def dialtion(im, k=1):
    element = cv.getStructuringElement(cv.MORPH_RECT, (2*k+1,2*k+1))
    im_dia = cv.dilate(im, element)
    return im_dia

#计算3D彩色网格空间像素分布密度
def density(pixels, bins=16):
    den = np.zeros((bins,bins,bins), np.int32)
    t = np.linspace(start=0, stop=256, num=bins+1, endpoint=True)
    ch_0 = pixels[:,0]
    ch_1 = pixels[:,1]
    ch_2 = pixels[:,2]
    #print(ch_b.shape, ch_g.shape, ch_r.shape)
    for b in range(bins):
        mask_b = np.bitwise_and(ch_0 >= t[b], ch_0 < t[b+1])
        for g in range(bins):
            mask_g = np.bitwise_and(ch_1 >= t[g], ch_1 < t[g+1])
            for r in range(bins):
                mask_r = np.bitwise_and(ch_2 >= t[r], ch_2 < t[r+1])
                mask = np.bitwise_and(mask_r, np.bitwise_and(mask_b, mask_g))
                den[b,g,r] = np.count_nonzero(mask)
                #if (den[b,g,r] > 0):
                #    print(den[b,g,r], (b,g,r))
    return den

#计算两个HSV颜色值之间的欧氏距离
def distance(c1, c2):
    dh = np.abs(int(c1[0]) - int(c2[0]))
    dh = 255 - dh if dh > 127 else dh
    ds = int(c1[1]) - int(c2[1])
    dv = int(c1[2]) - int(c2[2])
    return np.sqrt(dh * dh + ds * ds + dv * dv)


#根据3D彩色网格空间像素分布密度，查找除指定初始中心点外的另外2个初始中心点
#初始中心点用于k-means算法
def find_another_2center(density, first_center, min_distance):
    bins = density.shape
    c1 = (int(first_center[0] / bins[0] + 0.5),
          int(first_center[1] / bins[1] + 0.5),
          int(first_center[2] / bins[2] + 0.5))
    cur_max = 0
    c2 = (0,0,0)
    for b in range(bins[0]):
        for g in range(bins[1]):
            for r in range(bins[2]):
                if (distance((b,g,r), c1) > min_distance) and density[b,g,r] > cur_max:
                    cur_max = density[b,g,r]
                    c2 = (b,g,r)
    cur_max = 0
    c3 = (0,0,0)
    for b in range(bins[0]):
        for g in range(bins[1]):
            for r in range(bins[2]):
                if (distance((b,g,r), c1) > min_distance) and (distance((b,g,r), c2) > min_distance) and (density[b,g,r] > cur_max):
                    cur_max = density[b,g,r]
                    c3 = (b,g,r)
    k1 = 256 // bins[0]
    k2 = 256 // bins[1]
    k3 = 256 // bins[2]
    return (c2[0]*k1, c2[1]*k2, c2[2]*k3), (c3[0]*k1,c3[1]*k2,c3[2]*k3)

#根据指定的中心值初始化像素的label
#为k-means算法准备
def initial_label(pixels, c1, c2, c3):
    label = np.zeros((pixels.shape[0],), np.int32)
    dis = np.zeros((3,), np.int32)
    for i,v in enumerate(pixels):
        dis[0] = distance(c1, v)
        dis[1] = distance(c2, v)
        dis[2] = distance(c3, v)
        label[i] =np.argmin(dis)
    return label

#计算mask的空间分布离散度，返回值用于选择网格背景
def coef_of_var(mask, rows, cols):
    blk_h = mask.shape[0] // rows
    blk_w = mask.shape[1] // cols
    non_zeros = np.zeros(rows*cols, np.float32)
    for r in range(rows):
        for c in range(cols):
            rr = r * blk_h
            cc = c * blk_w
            non_zeros[r*cols+c] = np.count_nonzero(mask[rr:rr+blk_h, cc:cc+blk_w])
    return (non_zeros.std() / non_zeros.mean())


if __name__ == '__main__':
    fn = ""
    if len(sys.argv) < 2:
        fn = '0-input.png'
    else:
        fn = sys.argv[1]
    
    im = cv.imread(fn)
    sharped = unsharp_mask(im,k=1.5)
    cv.imwrite('1-sharped.png', sharped)

    bin = color2bin(sharped)
    cv.imwrite('2-bin.png', 255 - bin)
    
    dia_bin = dialtion(bin, 1)
    cv.imwrite('3-dialted.png', dia_bin)

    edge = np.bitwise_xor(bin, dia_bin)
    cv.imwrite('4-edge.png', edge)
    
    #将RGB转为HSV，减小像素通道之间的相关性，提高聚类算法的精度
    im_hsv = cv.cvtColor(im, cv.COLOR_BGR2HSV_FULL)

    #笔划和网格线外边缘的像素均值作为第一个初始中心点
    c1 = np.mean(im_hsv[edge>0], 0).astype(np.uint8).tolist()

    #提取参数聚类的像素坐标与像素值
    coords = np.nonzero(dia_bin > 0)
    pixels = im_hsv[dia_bin > 0]

    #计算3D彩色空间像素分布密度
    den = density(pixels)

    #在3D彩色空间像素分布密度中查找另外2个初始中心点
    c2, c3 = find_another_2center(den, c1, 2*np.sqrt(3))

    #初始化像素的分类
    label_ini = initial_label(pixels, c1, c2, c3)

    #执行kmean聚类
    max_iter = 7
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, max_iter, 1.0)
    ret,label,center = cv.kmeans(pixels.astype(np.float32), 3, label_ini, criteria, max_iter, cv.KMEANS_USE_INITIAL_LABELS)

    #根据聚类的label结果渲染为mask和原图上叠加的彩色
    clr_table = np.array([[255,0,0],[0,255,0],[0,0,255]], np.uint8)
    row = coords[0]
    col = coords[1]
    colorized = [im.copy(), im.copy(), im.copy()]
    masks = [np.zeros(im.shape, np.uint8), np.zeros(im.shape, np.uint8), np.zeros(im.shape, np.uint8)]
    overlay = im.copy()
    for i, v in enumerate(label):
        colorized[v][row[i], col[i]] = clr_table[v]
        overlay[row[i],col[i]] = clr_table[v]
        masks[v][row[i], col[i]] = 255

    #计算各label的空间分布离散系数
    v0 = coef_of_var(masks[0], 4, 7)
    v1 = coef_of_var(masks[1], 4, 7)
    v2 = coef_of_var(masks[2], 4, 7)
    vars = np.array([v0,v1,v2])

    #三类像素中，亮度最大的为背景，
    #剩下的两类像素，空间分布离散系数小的是笔迹
    vars[np.argmax(np.array(center)[:,2])] = 0
    index = np.argmax(v2)

    cv.imwrite('5-result.png', 255 - masks[index])
    cv.imwrite('6-colorized.png', colorized[index])
    cv.imwrite('7-overlay.png', overlay)

    print('initial center        = ', [c1, c2, c3])
    print('k-means output center = ', center.astype(np.uint8).tolist())

    #'''
    # 绘制像素散点图
    fig = plt.figure()
    ax = ax = fig.add_subplot(projection='3d')
    im_rgb = cv.cvtColor(im, cv.COLOR_BGR2RGB)
    colors = clr_table[2-label]/255.0
    ax.scatter(pixels[:,0], pixels[:,1], pixels[:,2],marker = '.',c = colors)
    ax.set_ylabel('S', fontdict={'size': 15, 'color': 'red'})
    ax.set_xlabel('H', fontdict={'size': 15, 'color': 'red'})
    ax.set_zlabel('V', fontdict={'size': 15, 'color': 'red'})
    plt.show()
    #'''

